import inspect
from collections import OrderedDict
from typing import Callable, Dict, List, Union

from lagent.actions.base_action import BaseAction
from lagent.actions.builtin_actions import FinishAction, InvalidAction, NoAction
from lagent.hooks import Hook, RemovableHandle
from lagent.schema import ActionReturn, ActionValidCode, AgentMessage, FunctionCall
from lagent.utils import create_object


class ActionExecutor:
    """
    动作执行器类
    
   负责管理和执行各种动作。它初始化动作列表，处理无效动作、空动作和完成动作，并支持hook函数的注册和执行。
    
    参数:
        actions: 动作或动作列表
        invalid_action: 无效动作处理器，默认为InvalidAction
        no_action: 空动作处理器，默认为NoAction
        finish_action: 完成动作处理器，默认为FinishAction
        finish_in_action: 是否将完成动作加入动作列表，默认False
        hooks: 钩子函数列表，用于处理动作执行前后的逻辑
    """

    def __init__(
        self,
        actions: Union[BaseAction, List[BaseAction], Dict, List[Dict]],
        invalid_action: BaseAction = dict(type=InvalidAction),
        no_action: BaseAction = dict(type=NoAction),
        finish_action: BaseAction = dict(type=FinishAction),
        finish_in_action: bool = False,
        hooks: List[Dict] = None,
    ):
        # 将单个动作转换为列表
        if not isinstance(actions, list):
            actions = [actions]
            
        # 创建完成动作实例
        finish_action = create_object(finish_action)
        if finish_in_action:
            actions.append(finish_action)
            
        # 初始化所有动作
        for i, action in enumerate(actions):
            actions[i] = create_object(action)
        self.actions = {action.name: action for action in actions}

        # 处理未注册或无效的动作请求(当Agent请求执行一个系统中不存在的动作)
        self.invalid_action = create_object(invalid_action)
        # 处理不需要执行任何具体操作的情况(当Agent决定当前不需要执行任何操作)
        self.no_action = create_object(no_action)
        # 处理任务完成的标记(当Agent认为当前任务已经完成)
        self.finish_action = finish_action
        
        # 初始化钩子系统
        self._hooks: Dict[int, Hook] = OrderedDict()
        if hooks:
            for hook in hooks:
                hook = create_object(hook)
                self.register_hook(hook)

    def description(self) -> List[Dict]:
        """
        获取所有动作的描述
        
        返回:
            List[Dict]: 动作描述列表
        """
        actions = []
        for action_name, action in self.actions.items():
            if action.is_toolkit:  # 处理工具包类型的动作
                for api in action.description['api_list']:
                    api_desc = api.copy()
                    api_desc['name'] = f"{action_name}.{api_desc['name']}"
                    actions.append(api_desc)
            else:  # 处理普通动作
                action_desc = action.description.copy()
                actions.append(action_desc)
        return actions

    def __contains__(self, name: str):
        """检查是否包含指定名称的动作"""
        return name in self.actions

    def keys(self):
        """获取所有动作的名称列表"""
        return list(self.actions.keys())

    def __setitem__(self, name: str, action: Union[BaseAction, Dict]):
        """添加或更新动作"""
        action = create_object(action)
        self.actions[action.name] = action

    def __delitem__(self, name: str):
        """删除指定动作"""
        del self.actions[name]

    def forward(self, name, parameters, **kwargs) -> ActionReturn:
        """
        执行动作的核心逻辑
        
        @description 根据动作名称执行相应的处理逻辑：
            1. 对于未注册动作：
               - no_action: 表示不需要执行任何动作的情况
               - finish_action: 表示完成当前任务的情况
               - 其他未注册动作: 作为无效动作处理
            2. 对于已注册动作：
               - 使用注册的处理器执行动作
               - 设置动作状态为开放
        
        @param {str} action_name - 要执行的动作名称
        @param {dict} parameters - 动作执行所需的参数
        @param {str} name - 动作的标识名称
        @param {str} api_name - API调用的名称
        
        @return {ActionReturn} action_return - 动作执行的结果对象，包含：
            - 执行状态
            - 执行结果
            - 错误信息（如果有）
            - 有效性标志
        """
        # 解析动作名称和API名称
        action_name, api_name = (
            name.split('.') if '.' in name else (name, 'run'))
        action_return: ActionReturn = ActionReturn()
        
        # 检查动作名称是否在已注册的动作列表中
        if action_name not in self:
            # 处理空动作：当AI决定不需要执行任何具体操作时
            if name == self.no_action.name:
                action_return = self.no_action(parameters)
            # 处理完成动作：当AI认为当前任务已完成时
            elif name == self.finish_action.name:
                action_return = self.finish_action(parameters)
            # 处理无效动作：当请求的动作未在系统中注册时
            else:
                action_return = self.invalid_action(parameters)
        else:
            # 执行已注册的动作处理器
            # parameters: 包含动作执行所需的具体参数
            # api_name: 用于标识具体调用的API方法
            action_return = self.actions[action_name](parameters, api_name)
            # 将动作状态设置为开放（OPEN），表示动作执行系统处于可用状态
            action_return.valid = ActionValidCode.OPEN
        # 返回动作执行的结果对象
        return action_return

    def __call__(self,
                 message: AgentMessage,
                 session_id=0,
                 **kwargs) -> AgentMessage:
        """
        处理动作消息
        
        参数:
            message: 代理消息
            session_id: 会话ID
            
        返回:
            AgentMessage: 处理后的消息
        """
        # 执行前置钩子
        for hook in self._hooks.values():
            result = hook.before_action(self, message, session_id)
            if result:
                message = result

        # 验证消息格式并提取动作信息
        assert isinstance(message.content, FunctionCall) or (
            isinstance(message.content, dict) and 'name' in message.content
            and 'parameters' in message.content)
            
        if isinstance(message.content, dict):
            name = message.content.get('name')
            parameters = message.content.get('parameters')
        else:
            name = message.content.name
            parameters = message.content.parameters

        # 执行动作
        response_message = self.forward(
            name=name, parameters=parameters, **kwargs)
            
        # 格式化响应消息
        if not isinstance(response_message, AgentMessage):
            response_message = AgentMessage(
                sender=self.__class__.__name__,
                content=response_message,
            )

        # 执行后置钩子
        for hook in self._hooks.values():
            result = hook.after_action(self, response_message, session_id)
            if result:
                response_message = result
        return response_message

    def register_hook(self, hook: Callable):
        """注册钩子函数"""
        handle = RemovableHandle(self._hooks)
        self._hooks[handle.id] = hook
        return handle


class AsyncActionExecutor(ActionExecutor):
    """
    异步动作执行器
    
    提供异步执行动作的能力，继承自ActionExecutor
    """

    async def forward(self, name, parameters, **kwargs) -> ActionReturn:
        """
        异步执行指定的动作
        
        参数:
            name: 动作名称
            parameters: 动作参数
            
        返回:
            ActionReturn: 动作执行结果
        """
        action_name, api_name = (
            name.split('.') if '.' in name else (name, 'run'))
        action_return: ActionReturn = ActionReturn()
        
        # 处理不同类型的动作
        if action_name not in self:
            if name == self.no_action.name:
                action_return = self.no_action(parameters)
            elif name == self.finish_action.name:
                action_return = self.finish_action(parameters)
            else:
                action_return = self.invalid_action(parameters)
        else:
            action = self.actions[action_name]
            # 根据动作类型选择同步或异步执行
            if inspect.iscoroutinefunction(action.__call__):
                action_return = await action(parameters, api_name)
            else:
                action_return = action(parameters, api_name)
            action_return.valid = ActionValidCode.OPEN
        return action_return

    async def __call__(self,
                       message: AgentMessage,
                       session_id=0,
                       **kwargs) -> AgentMessage:
        """
        异步处理动作消息
        
        参数:
            message: 代理消息
            session_id: 会话ID
            
        返回:
            AgentMessage: 处理后的消息
        """
        # 执行前置钩子（支持同步和异步）
        for hook in self._hooks.values():
            if inspect.iscoroutinefunction(hook.before_action):
                result = await hook.before_action(self, message, session_id)
            else:
                result = hook.before_action(self, message, session_id)
            if result:
                message = result

        # 验证消息格式并提取动作信息
        assert isinstance(message.content, FunctionCall) or (
            isinstance(message.content, dict) and 'name' in message.content
            and 'parameters' in message.content)
            
        if isinstance(message.content, dict):
            name = message.content.get('name')
            parameters = message.content.get('parameters')
        else:
            name = message.content.name
            parameters = message.content.parameters

        # 异步执行动作
        response_message = await self.forward(
            name=name, parameters=parameters, **kwargs)
            
        # 格式化响应消息
        if not isinstance(response_message, AgentMessage):
            response_message = AgentMessage(
                sender=self.__class__.__name__,
                content=response_message,
            )

        # 执行后置钩子（支持同步和异步）
        for hook in self._hooks.values():
            if inspect.iscoroutinefunction(hook.after_action):
                result = await hook.after_action(self, response_message,
                                               session_id)
            else:
                result = hook.after_action(self, response_message, session_id)
            if result:
                response_message = result
        return response_message
